This change is required because we use the ZeroCopyByteStream class now to read the initial message. ZeroCopyByteStream returns memoryview objects instead of bytes or bytearray. s.decode is only implemented for s ∈ {bytes, bytearray}.

According to the Python documentation, for bytes.decode and codecs.decode the later implementation may behave different in case of errors.

bytes.decode states that

errors controls how decoding errors are handled. If 'strict' (the default), a UnicodeError exception is raised.

codecs.decode states that

The default error handler is 'strict' meaning that decoding errors raise ValueError (or a more codec specific subclass, such as UnicodeDecodeError)

As decoding errors are unexpected and the specific Exception type that is thrown should not matter, I believe this change is safe. However, if there are concerns we can change this implementation to the following:

return bytes(s).decode("utf-8") if self.use_unicode else s

This alternative implementation will invoke a memory copy to copy the memoryview contents into a bytes object before decoding. Given that this call is only invoked on deserialization of the initial message, this memory copy is probably acceptable as it is a one time cost.

Happy to rename this function if another name is prefered

I think it is fine for now

This change is required to make mypy happy due to the following call in worker.py:

# Cleanup
# Reset task context to None. This is a guard code to avoid residual context when worker
# reuse.
TaskContext._setTaskContext(None)
BarrierTaskContext._setTaskContext(None)

It is unclear to my why this type check if failing only now since this code has been in worker.py for a long time.

I don't think adding runtime assertions is a good idea. to make mypy happy, maybe typing.cast(memoryview, init_info.udf_info)?

While typing.cast(memoryview, init_info.udf_info) would help mypy, it also means that init_info.udf_info is not guaranteed to be of type memoryview. If someone were to change the code in WorkerInitInfo.from_stream to return a different type for the NON_UDF case, this issue would not be caught and might lead to subtle run-time errors. My rationale for asserting here was to prevent such a situation.

However, I agree that an assert is probably not the solution here. I revisited the problem, and I think I have found a nice solution using a TypeGuard instead. In worker_message.py, the following TypeGuard has been added:

def is_non_udf_info(
    udf_info: UdfInfoType,
    eval_type: int,
) -> TypeGuard[memoryview]:
    """TypeGuard that narrows udf_info to memoryview when eval_type is NON_UDF."""
    return eval_type == PythonEvalType.NON_UDF

This can be used in worker.py instead of the if eval_type == PythonEvalType.NON_UDF:, which resolves the typing error. Additionally, it keeps the typing logic close to the source where udf_info is defined. Let's discuss if you prefer a different solution!

I still prefer we have a flag and by default use the original code path. replacing it directly would need some benchmark proof.